#include "wavelet.h"

//
// Convolve the rows and downsample the columns. Use half-point symmetrization.
//
int
conv_row_down_col (data_t *y, data_t const *x, int m, int n, 
                   filt_t const* f, int lf)
{
	int ny = (n + lf-1) / 2;

	for (int i = 0; i < m; ++i){
		for (int j = 0; j < ny; ++j)
		{
			y[i + m*j] = 0.0f;

			for (int k = 0; k < lf; ++k)
			{
				int xj = 2*j+1 - (lf-1) + k;
				if (xj < 0) xj = -xj-1;
				if (xj >= n) xj = n-1 - (xj-n);

				y[i + m*j] += x[i + m*xj] * f[lf-1-k];
			}
		}
	}

	return ny;
}

//
// Convolve the columns and downsample the rows. Use half-point symmetrization.
//
int
conv_col_down_row (data_t *a, data_t const *y, int m, int n, 
                   filt_t const* f, int lf)
{
	int ma = (m + lf-1) / 2;

	for (int i = 0; i < ma; ++i){
		for (int j = 0; j < n; ++j)
		{
			a[i + ma*j] = 0.0f;

			for (int k = 0; k < lf; ++k)
			{
				int yi =  2*i+1 - (lf-1) + k;
				if (yi < 0) yi = -yi-1;
				if (yi >= m) yi = m-1 - (yi-m);

				a[i + ma*j] += y[yi + m*j] * f[lf-1-k];
			}
		}
	}

	return ma;
}

//
// One level of forward wavelet decomposition
//
void
forward_dwt2 (data_t *a, data_t *h, data_t *v, data_t *d,
              data_t const *x, data_t *y, int m, int n,
              filt_t const *lod, filt_t const *hid, int lf)
{
	int ny = conv_row_down_col (y, x, m, n, lod, lf);

	conv_col_down_row (a, y, m, ny, lod, lf);
	conv_col_down_row (h, y, m, ny, hid, lf);

	conv_row_down_col (y, x, m, n, hid, lf);

	conv_col_down_row (v, y, m, ny, lod, lf);
	conv_col_down_row (d, y, m, ny, hid, lf);
}

void
forward_wavelet (data_t *c, int nc, data_t *x, int m, int n,
                   int L, int *S, filt_t const *lod, filt_t const *hid, int lf)
{
	//if (!S) nc = transform_sizes (S, m, n, L, lf);
	//if (!c) c = new data_t [nc];

	data_t *h, *v, *d, *tmp, *a, *y, *z;

	a = new data_t[m*n];
	y = new data_t[m*n];
	z = new data_t[m*n];

	
	for (int i = 0; i < m*n; ++i) z[i] = x[i];

	h = c + S[0]*S[1];
	for (int l = 1; l <= L; ++l){
		h += 3*S[0 + 2*l]*S[1 + 2*l];
	}

	for (int l = L; l >= 1; --l)
	{
		int ml = S[0 + 2*l], nl = S[1 + 2*l];

		h = h - 3*ml*nl;
		v = h + ml*nl;
		d = v + ml*nl;

		forward_dwt2 (a, h, v, d, z, y, m, n, lod, hid, lf);

		tmp = z; z = a; a = tmp;
		m = ml; n = nl;
	}

	for (int i = 0; i < S[0]*S[1]; ++i) c[i] = z[i];

	delete [] a;
	delete [] y;
	delete [] z;
}

int
conv_col_up_row (data_t *y, data_t const *a, int m, int n,
                 filt_t const *f, int lf)
{
	int my = 2*m-1 + lf-1;

	for (int i = 0; i < my; ++i){
		for (int j = 0; j < n; ++j)
		{
			for (int k = (i - (lf-1)) & 1; k < lf; k += 2){
				int ai = (i - (lf-1) + k) >> 1;
				if (ai >= 0 and ai < m)
					y[i + my*j] += a[ai + m*j] * f[lf-1-k];
			}
		}
	}

	return my;
}

int
conv_row_up_col (data_t *y, data_t const *x, int m, int n,
                 filt_t const *f, int lf)
{
	int ny = 2*n-1 + lf-1;

	for (int i = 0; i < m; ++i){
		for (int j = 0; j < ny; ++j)
		{
			for (int k = (j - (lf-1)) & 1; k < lf; k += 2){
				int xj = (j - (lf-1) + k) >> 1;
				if (xj >= 0 and xj < n)
					y[i + m*j] += x[i + m*xj] * f[lf-1-k];
			}
		}
	}

	return ny;
}

void
inverse_dwt2 (data_t *x,
              data_t *a, data_t *h, data_t *v, data_t *d,
              data_t *y, int m, int n,
              filt_t const *lor, filt_t const *hir, int lf)
{
	int mx = 2*m-1 + lf-1;

	for (int i = 0; i < mx*n; ++i) y[i] = 0.0f;

	int nx = 2*n-1 + lf-1;

	for (int i = 0; i < mx*nx; ++i) x[i] = 0.0f;

	conv_col_up_row (y, a, m, n, lor, lf);
	conv_col_up_row (y, h, m, n, hir, lf);
	conv_row_up_col (x, y, mx, n, lor, lf);

	for (int i = 0; i < mx*n; ++i) y[i] = 0.0f;

	conv_col_up_row (y, v, m, n, lor, lf);
	conv_col_up_row (y, d, m, n, hir, lf);
	conv_row_up_col (x, y, mx, n, hir, lf);
}

void
inverse_wavelet (data_t *x, int mx, int nx, data_t *c, int nc,
                   int L, int *S, filt_t const *lor, filt_t const *hir, int lf)
{
	data_t *a, *h, *v, *d, *y, *z;

	//mx = S[0 + 2*(L+1)], nx = S[1 + 2*(L+1)];
	//if (not x) x = new data_t[mx * nx];

	int ms = S[0 + 2*L], ns = S[1 + 2*L];

	ms = 2*ms-1 + lf-1;
	ns = 2*ns-1 + lf-1;

	a = new data_t[ms * ns];
	y = new data_t[ms * ns];
	z = new data_t[ms * ns];

	int m = S[0], n = S[1];
	for (int i = 0; i < m*n; ++i) a[i] = c[i];
	h = c + m*n;

	for (int l = 1; l < L+1; ++l)
	{
		m = S[0 + 2*l];
		n = S[1 + 2*l];

		v = h + m*n;
		d = v + m*n;

		inverse_dwt2 (z, a, h, v, d, y, m, n, lor, hir, lf);

		int om = (int) floor(((2*m-1 + lf-1) - S[0 + 2*(l+1)]) / 2.0);
		int on = (int) floor(((2*n-1 + lf-1) - S[1 + 2*(l+1)]) / 2.0);

		for (int j = 0; j < S[1 + 2*(l+1)]; ++j){
			for (int i = 0; i < S[0 + 2*(l+1)]; ++i){
				a[i + S[2*(l+1)]*j] = z[om+i + (2*m-1+lf-1)*(on+j)];
			}
		}

		
		h += 3*m*n;
	}

	for (int i = 0; i < mx*nx; ++i) x[i] = a[i]; 

	delete [] y;
	delete [] z;
	delete [] a;
}

